import os
import random

import numpy as np
import torch

from util.seed import set_seed
from raw.get_model import get_model

def main(name):
    # 设置随机数种子
    set_seed()

    # 创建网络
    model = get_model(name)
    # 加载已保存的网络
    model.load_raw_state()
    # 加载测试数据
    test_loader = model.get_test_loader(batch_size = 1, shuffle = False)

    # 获取识别成功的数据
    model.eval()
    corr = []
    with torch.no_grad():
        for image, label in test_loader:
            image = image.to(model.device)
            label = label.to(model.device)
            output = model(image)
            pred = torch.argmax(output, dim = 1)
            if pred == label:
                llc = torch.argmin(output, dim = 1) # 找出概率最小的标签
                corr.append((image, label, llc))

        label_range = model(corr[0][0]).size(1) # 获取标签范围

    # 随机选取部分识别成功的数据
    candidates = random.sample(corr, 1000)

    candidate_images = []
    candidate_labels = []
    candidate_LLCs = []
    candidate_targets = []
    
    '''
     四类分别是图像、标签、least like class 、选中的目标
     ind是共选1000个image作为candidate中作为计数的变量，0~2是分别是image、label、llc
     在predict label后的结果的本身的label相同加入至corr当中从而可以得到
    '''

    for ind in range(len(candidates)):
        image = candidates[ind][0].cpu().numpy()[0]
        candidate_images.append(image)

        label = candidates[ind][1].cpu().numpy()[0]
        candidate_labels.append(label)

        llc = candidates[ind][2].cpu().numpy()[0]
        candidate_LLCs.append(llc)

        # 选择一个错误标签作为攻击的“误导目标”
        target = random.choice([i for i in range(label_range) if i != label])
        candidate_targets.append(target)

    # 创建保存目录
    outer_dir = 'data/candidate_data/'
    base_dir = outer_dir + name + '/'
    if name not in os.listdir(outer_dir):
        os.mkdir(base_dir)

    np.save(base_dir + 'images.npy', np.array(candidate_images))
    np.save(base_dir + 'labels.npy', np.array(candidate_labels))
    np.save(base_dir + 'llcs.npy', np.array(candidate_LLCs))
    np.save(base_dir + 'targets.npy', np.array(candidate_targets))
    print('Selection saved to ' + base_dir + '*.npy')
